from tempfile import mkstemp
from os import remove
from os import close

from zenpdb.Contacts import ACTAtomContacts
from zenpdb.Surface import ContactAtomSurface, AreaimolAtomSurface, ASUAreaimolAtomSurface, asu_contacts
from zenpdb.SecondaryStructure import SecondaryStructure
from zenpdb.PDBParser import PDBParser
from zenpdb.EntityManipulate import clean_ical, transfer_nearest_xtra, find_coordinations
from zenpdb.PDBWriter import PDBWriter
from zenpdb.BeQu import BeQu
from zenpdb.EntityMisc import einput
from zenpdb.Cluster import AttrResidueClusters
inp = 'inp.pdb'
oup = 'test.pdb'

pdb_parser = PDBParser()
pdb_writer = PDBWriter()
(fd, pdb_tmp) = mkstemp('.pdb')
close(fd)               # explicit close required


def process_file(input_file):
    # parse PDB file
    try:
        pdb_parser.set_file(input_file)
    except IOError:
        return 1

    # clean up structure
    try:
        structure = pdb_parser.get_structure(input_file[-8:-4]) # assume path/XXXX.pdb
        clean_ical(structure, pretend =False)
        structure.remove_hydrogens()
        structure.remove_hetero()
        structure.remove_empty()
    except:
        return 2

    # write temporary clean file
    try:
        pdb_writer.set_structure(structure)
        pdb_writer.save(pdb_tmp)
    except:
        return 3

    # prepare
    model = structure[(0,)]
    atoms = einput(model, 'A')
    residues = einput(model, 'R')
    residues.dispatch('set_coords')
    chains = einput(model, 'C')

    # run externals
        # distance contact
    ACTAtomContacts(pdb_tmp, model, contact ='inter 0.0 5.0')
        # surface contact in crystal
    ContactAtomSurface(pdb_tmp, model)
        # surface contact in asu
    ASUAreaimolAtomSurface(model)
    remove(pdb_tmp)

    # derived data
    asu_contacts(model)
    find_coordinations(model, 'A', cnt_key ='CNT_ACT_X', xtra_key = 'CN_ACT_X')

    # propagations
    model._data_propagate('merge','A','CNT_ACT_X', xtra =True)
    model._data_propagate('unique','A','CN_ACT_X', xtra =True)
    model._data_propagate('sum','A','CNT_AREAIMOL_ASA', xtra =True)
    model._data_propagate('sum','A','ASU_CNT_AREAIMOL_ASA', xtra =True)

    # transfer coordinations
    asar = residues._select_children(0.0, 'lt', 'CNT_AREAIMOL_ASA', xtra =True).values()
    cntr = residues._select_children({}, 'lt', 'CNT_ACT_X', xtra =True).values()

    transfer_nearest_xtra(asar, cntr, 'R', 'CN_ACT_X')
    AttrResidueClusters(asar, xtra_key ='CNT_CLUST', attr ='CN_ACT_X', xtra =True)
    BeQu('new.pdb', structure, 'R', b_key ='CNT_CLUST')
    return model